# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度
选用两种策略：
先采用 要么最优，要么最差
如果效果不好就选第二种
动作要么选最优的，要么选次优的

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
np.random.seed(1)
T = 2*np.pi
dt = np.pi/5
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex) 
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
action_space = np.array([0,1,2,3])#,5,6,7,8])

theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数
#总点数为 theta_num * varphi_num + 2(布洛赫球两极)

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set
#----------------------------------------------------------------------------------------------------
#动作直接选最优的
def step(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        # del action_list[fid_list.index(max(fid_list))]
        # del psi_list[fid_list.index(max(fid_list))]
        # del fid_list[fid_list.index(max(fid_list))]
        
        # best_action = fid_list.index(max(fid_list))
        # best_fid = max(fid_list)
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        del action_list[fid_list.index(max(fid_list))]
        del psi_list[fid_list.index(max(fid_list))]
        del fid_list[fid_list.index(max(fid_list))]
        
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#将测试集的保真度从小到大排列出来，来展示保真度分布
def sort_fid(test_fidelity_list):
    sort_fid = []
    for i in range (test_fidelity_list.shape[0]):
        b = test_fidelity_list[i,:]
        sort_fid  = np.append(sort_fid,b)
    sort_fid.sort()
    return sort_fid
#--------------------------------------------------------------------------------


def job(target_psi):
    fid_list = []
    for psi1 in init_set:
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        fid_max = 0
        fid_max1 = 0
        fid_max2 = 0
        fid_max0 = 0
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2)
        if fid_max < 0.99:
            step_n = 0
            F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
            psi = psi1
            while True:
                action, F, psi_ = step(psi,target_psi,F)
                fid_max0 = max(F,fid_max0)
                psi = psi_
                step_n += 1
                if fid_max0>0.999 or step_n>step_max:
                    break 
            fid_max = max(fid_max,fid_max0)  
        fid_list.append(fid_max)
    return  np.mean(fid_list)

def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list
    

if __name__ == '__main__':
    target_set = psi_set()
    # target_set = target_set[68:73]+target_set[89:94]# A area 0.9018504159122663
    target_set = target_set[36:40]+target_set[57:61]# B area 0.8862742472629985
    init_set = psi_set()
    # print(target_set)
    time1 = time()
    F_list = multicore()
    print(F_list)
    time2 = time()
    print(np.mean(F_list))
    print('time_cost is: ',time2-time1)
    
    #F_list.sort()# jiang liebiao an shunxu pailie.

#目标点集
# [matrix([[0.92387953+0.j],
#          [0.38268343+0.j]]),
#  matrix([[0.92387953+0.j        ],
#          [0.27059805+0.27059805j]]),
#  matrix([[9.23879533e-01+0.j        ],
#          [2.34326020e-17+0.38268343j]]),
#  matrix([[ 0.92387953+0.j        ],
#          [-0.27059805+0.27059805j]]),
#  matrix([[ 0.92387953+0.00000000e+00j],
#          [-0.38268343+4.68652041e-17j]]),
#  matrix([[ 0.92387953+0.j        ],
#          [-0.27059805-0.27059805j]]),
#  matrix([[ 9.23879533e-01+0.j        ],
#          [-7.02978061e-17-0.38268343j]]),
#  matrix([[0.92387953+0.j        ],
#          [0.27059805-0.27059805j]]),
#  matrix([[0.70710678+0.j],
#          [0.70710678+0.j]]),
#  matrix([[0.70710678+0.j ],
#          [0.5       +0.5j]]),
#  matrix([[7.07106781e-01+0.j        ],
#          [4.32978028e-17+0.70710678j]]),
#  matrix([[ 0.70710678+0.j ],
#          [-0.5       +0.5j]]),
#  matrix([[ 0.70710678+0.00000000e+00j],
#          [-0.70710678+8.65956056e-17j]]),
#  matrix([[ 0.70710678+0.j ],
#          [-0.5       -0.5j]]),
#  matrix([[ 7.07106781e-01+0.j        ],
#          [-1.29893408e-16-0.70710678j]]),
#  matrix([[0.70710678+0.j ],
#          [0.5       -0.5j]]),
#  matrix([[0.38268343+0.j],
#          [0.92387953+0.j]]),
#  matrix([[0.38268343+0.j        ],
#          [0.65328148+0.65328148j]]),
#  matrix([[3.82683432e-01+0.j        ],
#          [5.65713056e-17+0.92387953j]]),
#  matrix([[ 0.38268343+0.j        ],
#          [-0.65328148+0.65328148j]]),
#  matrix([[ 0.38268343+0.00000000e+00j],
#          [-0.92387953+1.13142611e-16j]]),
#  matrix([[ 0.38268343+0.j        ],
#          [-0.65328148-0.65328148j]]),
#  matrix([[ 3.82683432e-01+0.j        ],
#          [-1.69713917e-16-0.92387953j]]),
#  matrix([[0.38268343+0.j        ],
#          [0.65328148-0.65328148j]]),
#  matrix([[1.+0.j],
#          [0.+0.j]]),
#  matrix([[0.+0.j],
#          [1.+0.j]])]

#对应的保真度
# [0.9983698705101789 0.9985297684685446 0.9965556731244701 0.9955035177482916 0.994659297747708  0.995244267478151 0.9948305958230452 0.9932993913448795 0.9915782021447928 0.9924676880094832  0.9875100370670282 0.9863550322317771 0.9781222755123549 0.9564192946866689 0.9386778808982916 0.954011041740779 0.9829874929761757 0.9897813681414824 0.9940381621106791 0.9971115247924699 0.9975151327355625
  
  
#   0.9938934766086691 0.9973600056420892 0.9901710237625261 0.9890858242655995 0.9928267963250814 0.9880732980256026 0.9905488513787031 0.9903229788360068 0.990013665733689  0.9879469284783713 0.9919367429057087 0.9931635004756802 0.9866725840155094 0.9679130999500374  0.942501122724321 0.9135883702283729 0.8614099677128284 0.8028664867284254 0.8050867145043826  0.9915518051024743 0.9943030276416935
  
  
#   0.996878023881857 0.9959114385162455 0.9961067318812273 0.9917896705872514 0.9832462282117398 0.9723382875522987 0.9751390649984744 0.9747268921739078 0.9751195495226357 0.9750769417320411 0.9797814534164084 0.9944686022185072 0.9909133010119859 0.9867483422335799 0.9741638695649669 0.9523262300317582 0.9191505208419719 0.9145908143554086 0.9211748737008398 0.9402208411809976 0.9855212663460462 
  
  
#   0.9858414724802377 0.9955750647595535 0.9889727802621853 0.9806549800856832 0.9663421131510306 0.9345583146736149 0.9127178165464349 0.9166525420148361 0.9266097510553934 0.9542967546627172 0.9954706502846927 0.9968525526946006 0.996375760236706 0.9945645920421089 0.9887744348909606 0.9783443778942091 0.9712503451885746 0.9749152340266849 0.9768194861334132 0.9733283588188459 0.9729287447314083 
  

#   0.9937273011869621  0.9912344322902487 0.9829579366661805 0.950900792052235 0.9286408850871837 0.886104620700943 0.8319793129680799 0.7957292923461645 0.8656387227181339 0.9942170314363468 0.9950178530740493 0.9974105570751686 0.994385570776569 0.988638818503053 0.9916635522270711 0.9906551527586371 0.9879055548526442 0.991430297399177 0.9900675390465744 0.9889882720652874 0.9903879222325596 
  
  
#   0.9878005668318692 0.984728443919675 0.9692615853858224 0.9436409673870634 0.9477430808721119 0.971795064394405 0.9911357637520948 0.9933472589480354 0.996361391335465 0.9960969299975431 0.9977650937215402 0.9980512278508527 0.9973555664168654 0.9953965906731916 0.9942512406245512 0.9951728110097572 0.9946988960873684 0.994796202730763 0.9919389623911052 0.9930110516765831 0.9912947461197672 
  
  
#   0.9917313651779966, 
  
  
#   0.9911923226825154]

# 0.9725647765967409

# time_cost is:  15.473009824752808

# #保真度排序
# #F_list.sort()# 将列表按顺序排列.
# #print(F_list)

# [0.7957292923461645, 0.8028664867284254, 0.8050867145043826, 0.8319793129680799, 0.8614099677128284, 
#  0.8656387227181339, 0.886104620700943, 0.9127178165464349, 0.9135883702283729, 0.9145908143554086, 
#  0.9166525420148361, 0.9191505208419719, 0.9211748737008398, 0.9266097510553934, 0.9286408850871837, 
#  0.9345583146736149, 0.9386778808982916, 0.9402208411809976, 0.942501122724321, 0.9436409673870634, 
#  0.9477430808721119, 0.950900792052235, 0.9523262300317582, 0.954011041740779, 0.9542967546627172, 
#  0.9564192946866689, 0.9663421131510306, 0.9679130999500374, 0.9692615853858224, 0.9712503451885746, 
#  0.971795064394405, 0.9723382875522987, 0.9729287447314083, 0.9733283588188459, 0.9741638695649669, 
#  0.9747268921739078, 0.9749152340266849, 0.9750769417320411, 0.9751195495226357, 0.9751390649984744, 
#  0.9768194861334132, 0.9781222755123549, 0.9783443778942091, 0.9797814534164084, 0.9806549800856832, 
#  0.9829579366661805, 0.9829874929761757, 0.9832462282117398, 0.984728443919675, 0.9855212663460462, 
#  0.9858414724802377, 0.9863550322317771, 0.9866725840155094, 0.9867483422335799, 0.9875100370670282, 
#  0.9878005668318692, 0.9879055548526442, 0.9879469284783713, 0.9880732980256026, 0.988638818503053, 
#  0.9887744348909606, 0.9889727802621853, 0.9889882720652874, 0.9890858242655995, 0.9897813681414824, 
#  0.990013665733689, 0.9900675390465744, 0.9901710237625261, 0.9903229788360068, 0.9903879222325596, 
#  0.9905488513787031, 0.9906551527586371, 0.9909133010119859, 0.9911357637520948, 0.9911923226825154, 
#  0.9912344322902487, 0.9912947461197672, 0.991430297399177, 0.9915518051024743, 0.9915782021447928, 
#  0.9916635522270711, 0.9917313651779966, 0.9917896705872514, 0.9919367429057087, 0.9919389623911052, 
#  0.9924676880094832, 0.9928267963250814, 0.9930110516765831, 0.9931635004756802, 0.9932993913448795, 
#  0.9933472589480354, 0.9937273011869621, 0.9938934766086691, 0.9940381621106791, 0.9942170314363468, 
#  0.9942512406245512, 0.9943030276416935, 0.994385570776569, 0.9944686022185072, 0.9945645920421089, 
#  0.994659297747708, 0.9946988960873684, 0.994796202730763, 0.9948305958230452, 0.9950178530740493, 
#  0.9951728110097572, 0.995244267478151, 0.9953965906731916, 0.9954706502846927, 0.9955035177482916, 
#  0.9955750647595535, 0.9959114385162455, 0.9960969299975431, 0.9961067318812273, 0.996361391335465, 
#  0.996375760236706, 0.9965556731244701, 0.9968525526946006, 0.996878023881857, 0.9971115247924699, 
#  0.9973555664168654, 0.9973600056420892, 0.9974105570751686, 0.9975151327355625, 0.9977650937215402, 
#  0.9980512278508527, 0.9983698705101789, 0.9985297684685446]



